"""
电磁场预测模型损失函数
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from .physics_equations import MHDEquations, MaxwellEquations


class LpLoss(nn.Module):
    """Lp损失函数"""
    
    def __init__(self, p=2):
        super(LpLoss, self).__init__()
        """初始化Lp损失函数
        
        参数:
            p: p范数参数，默认为2
        """
        self.p = p
    
    def forward(self, pred, true):
        """计算Lp损失
        
        参数:
            pred: 预测值
            true: 真实值
            
        返回:
            Lp损失值
        """
        diff = pred - true
        loss = torch.mean(torch.abs(diff) ** self.p)
        return loss ** (1 / self.p)


class MHDLoss(nn.Module):
    """MHD方程损失函数"""
    
    def __init__(
        self,
        nu=1e-4,
        eta=1e-4,
        rho0=1.0,
        data_weight=1.0,
        ic_weight=1.0,
        pde_weight=1.0,
        constraint_weight=1.0,
        use_data_loss=True,
        use_ic_loss=True,
        use_pde_loss=True,
        use_constraint_loss=True,
        u_weight=1.0,
        v_weight=1.0,
        Bx_weight=1.0,
        By_weight=1.0,
        div_B_weight=1.0,
        div_vel_weight=1.0,
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        use_weighted_mean=False
    ):
        super(MHDLoss, self).__init__()
        """初始化MHD损失函数
        
        参数:
            nu: 运动粘度
            eta: 磁扩散率
            rho0: 参考密度
            data_weight: 数据损失权重
            ic_weight: 初始条件损失权重
            pde_weight: PDE损失权重
            constraint_weight: 约束损失权重
            use_data_loss: 是否使用数据损失
            use_ic_loss: 是否使用初始条件损失
            use_pde_loss: 是否使用PDE损失
            use_constraint_loss: 是否使用约束损失
            u_weight: u分量权重
            v_weight: v分量权重
            Bx_weight: Bx分量权重
            By_weight: By分量权重
            div_B_weight: 磁场散度约束权重
            div_vel_weight: 速度散度约束权重
            Lx, Ly, Lz: 域尺寸
            use_weighted_mean: 是否使用加权平均
        """
        self.nu = nu
        self.eta = eta
        self.rho0 = rho0
        self.data_weight = data_weight
        self.ic_weight = ic_weight
        self.pde_weight = pde_weight
        self.constraint_weight = constraint_weight
        self.use_data_loss = use_data_loss
        self.use_ic_loss = use_ic_loss
        self.use_pde_loss = use_pde_loss
        self.use_constraint_loss = use_constraint_loss
        self.u_weight = u_weight
        self.v_weight = v_weight
        self.Bx_weight = Bx_weight
        self.By_weight = By_weight
        self.div_B_weight = div_B_weight
        self.div_vel_weight = div_vel_weight
        self.Lx = Lx
        self.Ly = Ly
        self.Lz = Lz
        self.use_weighted_mean = use_weighted_mean
        
        # 初始化MHD方程
        self.mhd_eq = MHDEquations(nu, eta, rho0)
        
        # 初始化损失函数
        self.lp_loss = LpLoss(p=2)
        
        # 空间步长
        self.dx = Lx / (pred.shape[-3] - 1) if hasattr(self, 'pred') else 1.0
        self.dy = Ly / (pred.shape[-2] - 1) if hasattr(self, 'pred') else 1.0
        self.dz = Lz / (pred.shape[-1] - 1) if hasattr(self, 'pred') else 1.0
        
        if not self.use_data_loss:
            self.data_weight = 0
        if not self.use_ic_loss:
            self.ic_weight = 0
        if not self.use_pde_loss:
            self.pde_weight = 0
        if not self.use_constraint_loss:
            self.constraint_weight = 0
    
    def forward(self, pred, true, inputs, return_loss_dict=False):
        """计算MHD损失
        
        参数:
            pred: 预测值，形状为 [batch_size, x_dim, y_dim, z_dim, 4]
            true: 真实值，形状为 [batch_size, x_dim, y_dim, z_dim, 4]
            inputs: 输入值，形状为 [batch_size, x_dim, y_dim, z_dim, 7]
            return_loss_dict: 是否返回详细损失字典
            
        返回:
            总损失值，或(总损失值, 损失字典)
        """
        # 更新空间步长
        self.dx = self.Lx / (pred.shape[-3] - 1)
        self.dy = self.Ly / (pred.shape[-2] - 1)
        self.dz = self.Lz / (pred.shape[-1] - 1)
        
        if not return_loss_dict:
            loss = self.compute_loss(pred, true, inputs)
            return loss
        else:
            loss, loss_dict = self.compute_losses(pred, true, inputs)
            return loss, loss_dict
    
    def compute_loss(self, pred, true, inputs):
        """计算加权损失
        
        参数:
            pred: 预测值
            true: 真实值
            inputs: 输入值
            
        返回:
            总损失值
        """
        pred = pred.reshape(true.shape)
        
        # 提取分量
        u = pred[..., 0]
        v = pred[..., 1]
        Bx = pred[..., 2]
        By = pred[..., 3]
        
        # 构建速度场和磁场
        velocity = torch.stack([u, v, torch.zeros_like(u)], dim=-1)
        magnetic_field = torch.stack([Bx, By, torch.zeros_like(Bx)], dim=-1)
        
        # 数据损失
        if self.use_data_loss:
            loss_data = self.data_loss(pred, true)
        else:
            loss_data = 0
        
        # 初始条件损失
        if self.use_ic_loss:
            loss_ic = self.ic_loss(pred, inputs)
        else:
            loss_ic = 0
        
        # PDE损失
        if self.use_pde_loss:
            Du, Dv, DBx, DBy = self.mhd_residual(u, v, Bx, By, velocity, magnetic_field)
            loss_pde = self.pde_loss(Du, Dv, DBx, DBy)
        else:
            loss_pde = 0
        
        # 约束损失
        if self.use_constraint_loss:
            div_vel, div_B = self.mhd_eq.divergence_constraint(velocity, magnetic_field, self.dx, self.dy, self.dz)
            loss_constraint = self.constraint_loss(div_vel, div_B)
        else:
            loss_constraint = 0
        
        # 计算权重总和
        if self.use_weighted_mean:
            weight_sum = (
                self.data_weight
                + self.ic_weight
                + self.pde_weight
                + self.constraint_weight
            )
        else:
            weight_sum = 1.0
        
        # 总损失
        loss = (
            self.data_weight * loss_data
            + self.ic_weight * loss_ic
            + self.pde_weight * loss_pde
            + self.constraint_weight * loss_constraint
        ) / weight_sum
        
        return loss
    
    def compute_losses(self, pred, true, inputs):
        """计算加权损失和详细损失字典
        
        参数:
            pred: 预测值
            true: 真实值
            inputs: 输入值
            
        返回:
            总损失值和损失字典
        """
        pred = pred.reshape(true.shape)
        
        # 提取分量
        u = pred[..., 0]
        v = pred[..., 1]
        Bx = pred[..., 2]
        By = pred[..., 3]
        
        # 构建速度场和磁场
        velocity = torch.stack([u, v, torch.zeros_like(u)], dim=-1)
        magnetic_field = torch.stack([Bx, By, torch.zeros_like(Bx)], dim=-1)
        
        loss_dict = {}
        
        # 数据损失
        if self.use_data_loss:
            loss_data, loss_u, loss_v, loss_Bx, loss_By = self.data_loss(
                pred, true, return_all_losses=True
            )
            loss_dict["loss_data"] = loss_data
            loss_dict["loss_u"] = loss_u
            loss_dict["loss_v"] = loss_v
            loss_dict["loss_Bx"] = loss_Bx
            loss_dict["loss_By"] = loss_By
        else:
            loss_data = 0
        
        # 初始条件损失
        if self.use_ic_loss:
            loss_ic, loss_u_ic, loss_v_ic, loss_Bx_ic, loss_By_ic = self.ic_loss(
                pred, inputs, return_all_losses=True
            )
            loss_dict["loss_ic"] = loss_ic
            loss_dict["loss_u_ic"] = loss_u_ic
            loss_dict["loss_v_ic"] = loss_v_ic
            loss_dict["loss_Bx_ic"] = loss_Bx_ic
            loss_dict["loss_By_ic"] = loss_By_ic
        else:
            loss_ic = 0
        
        # PDE损失
        if self.use_pde_loss:
            Du, Dv, DBx, DBy = self.mhd_residual(u, v, Bx, By, velocity, magnetic_field)
            loss_pde, loss_Du, loss_Dv, loss_DBx, loss_DBy = self.pde_loss(
                Du, Dv, DBx, DBy, return_all_losses=True
            )
            loss_dict["loss_pde"] = loss_pde
            loss_dict["loss_Du"] = loss_Du
            loss_dict["loss_Dv"] = loss_Dv
            loss_dict["loss_DBx"] = loss_DBx
            loss_dict["loss_DBy"] = loss_DBy
        else:
            loss_pde = 0
        
        # 约束损失
        if self.use_constraint_loss:
            div_vel, div_B = self.mhd_eq.divergence_constraint(velocity, magnetic_field, self.dx, self.dy, self.dz)
            loss_constraint, loss_div_vel, loss_div_B = self.constraint_loss(
                div_vel, div_B, return_all_losses=True
            )
            loss_dict["loss_constraint"] = loss_constraint
            loss_dict["loss_div_vel"] = loss_div_vel
            loss_dict["loss_div_B"] = loss_div_B
        else:
            loss_constraint = 0
        
        # 计算权重总和
        if self.use_weighted_mean:
            weight_sum = (
                self.data_weight
                + self.ic_weight
                + self.pde_weight
                + self.constraint_weight
            )
        else:
            weight_sum = 1.0
        
        # 总损失
        loss = (
            self.data_weight * loss_data
            + self.ic_weight * loss_ic
            + self.pde_weight * loss_pde
            + self.constraint_weight * loss_constraint
        ) / weight_sum
        
        return loss, loss_dict
    
    def data_loss(self, pred, true, return_all_losses=False):
        """计算数据损失
        
        参数:
            pred: 预测值
            true: 真实值
            return_all_losses: 是否返回各分量损失
            
        返回:
            数据损失值，或(总损失, 各分量损失)
        """
        if return_all_losses:
            loss_u = self.lp_loss(pred[..., 0], true[..., 0])
            loss_v = self.lp_loss(pred[..., 1], true[..., 1])
            loss_Bx = self.lp_loss(pred[..., 2], true[..., 2])
            loss_By = self.lp_loss(pred[..., 3], true[..., 3])
            
            loss = (
                self.u_weight * loss_u
                + self.v_weight * loss_v
                + self.Bx_weight * loss_Bx
                + self.By_weight * loss_By
            )
            
            return loss, loss_u, loss_v, loss_Bx, loss_By
        else:
            loss = self.lp_loss(pred, true)
            return loss
    
    def ic_loss(self, pred, inputs, return_all_losses=False):
        """计算初始条件损失
        
        参数:
            pred: 预测值
            inputs: 输入值
            return_all_losses: 是否返回各分量损失
            
        返回:
            初始条件损失值，或(总损失, 各分量损失)
        """
        # 获取初始时刻的预测值和真实值
        pred_ic = pred[:, 0:1, ...]  # 只取第一个时间步
        true_ic = inputs[..., -4:]  # 输入的最后4个维度是初始条件
        
        if return_all_losses:
            loss_u = self.lp_loss(pred_ic[..., 0], true_ic[..., 0])
            loss_v = self.lp_loss(pred_ic[..., 1], true_ic[..., 1])
            loss_Bx = self.lp_loss(pred_ic[..., 2], true_ic[..., 2])
            loss_By = self.lp_loss(pred_ic[..., 3], true_ic[..., 3])
            
            loss = (
                self.u_weight * loss_u
                + self.v_weight * loss_v
                + self.Bx_weight * loss_Bx
                + self.By_weight * loss_By
            )
            
            return loss, loss_u, loss_v, loss_Bx, loss_By
        else:
            loss = self.lp_loss(pred_ic, true_ic)
            return loss
    
    def mhd_residual(self, u, v, Bx, By, velocity, magnetic_field):
        """计算MHD方程残差
        
        参数:
            u, v: 速度分量
            Bx, By: 磁场分量
            velocity: 速度场
            magnetic_field: 磁场
            
        返回:
            MHD方程残差
        """
        # 计算时间导数
        du_dt = torch.gradient(u, dim=1)[0]  # 假设第一维是时间
        dv_dt = torch.gradient(v, dim=1)[0]
        dBx_dt = torch.gradient(Bx, dim=1)[0]
        dBy_dt = torch.gradient(By, dim=1)[0]
        
        # 计算速度方程残差
        dv_dt_pred = self.mhd_eq.velocity_equation(
            velocity, magnetic_field, dx=self.dx, dy=self.dy, dz=self.dz
        )
        Du = du_dt - dv_dt_pred[..., 0]
        Dv = dv_dt - dv_dt_pred[..., 1]
        
        # 计算磁场方程残差
        dB_dt_pred = self.mhd_eq.magnetic_equation(
            velocity, magnetic_field, dx=self.dx, dy=self.dy, dz=self.dz
        )
        DBx = dBx_dt - dB_dt_pred[..., 0]
        DBy = dBy_dt - dB_dt_pred[..., 1]
        
        return Du, Dv, DBx, DBy
    
    def pde_loss(self, Du, Dv, DBx, DBy, return_all_losses=False):
        """计算PDE损失
        
        参数:
            Du, Dv, DBx, DBy: MHD方程残差
            return_all_losses: 是否返回各分量损失
            
        返回:
            PDE损失值，或(总损失, 各分量损失)
        """
        if return_all_losses:
            loss_Du = self.lp_loss(Du, torch.zeros_like(Du))
            loss_Dv = self.lp_loss(Dv, torch.zeros_like(Dv))
            loss_DBx = self.lp_loss(DBx, torch.zeros_like(DBx))
            loss_DBy = self.lp_loss(DBy, torch.zeros_like(DBy))
            
            loss = (
                self.Du_weight * loss_Du
                + self.Dv_weight * loss_Dv
                + self.DBx_weight * loss_DBx
                + self.DBy_weight * loss_DBy
            )
            
            return loss, loss_Du, loss_Dv, loss_DBx, loss_DBy
        else:
            loss = (
                self.Du_weight * self.lp_loss(Du, torch.zeros_like(Du))
                + self.Dv_weight * self.lp_loss(Dv, torch.zeros_like(Dv))
                + self.DBx_weight * self.lp_loss(DBx, torch.zeros_like(DBx))
                + self.DBy_weight * self.lp_loss(DBy, torch.zeros_like(DBy))
            )
            return loss
    
    def constraint_loss(self, div_vel, div_B, return_all_losses=False):
        """计算约束损失
        
        参数:
            div_vel: 速度散度
            div_B: 磁场散度
            return_all_losses: 是否返回各分量损失
            
        返回:
            约束损失值，或(总损失, 各分量损失)
        """
        if return_all_losses:
            loss_div_vel = self.lp_loss(div_vel, torch.zeros_like(div_vel))
            loss_div_B = self.lp_loss(div_B, torch.zeros_like(div_B))
            
            loss = (
                self.div_vel_weight * loss_div_vel
                + self.div_B_weight * loss_div_B
            )
            
            return loss, loss_div_vel, loss_div_B
        else:
            loss = (
                self.div_vel_weight * self.lp_loss(div_vel, torch.zeros_like(div_vel))
                + self.div_B_weight * self.lp_loss(div_B, torch.zeros_like(div_B))
            )
            return loss


class MaxwellLoss(nn.Module):
    """麦克斯韦方程损失函数"""
    
    def __init__(
        self,
        eps0=1.0,
        mu0=1.0,
        sigma=0.0,
        data_weight=1.0,
        pde_weight=1.0,
        constraint_weight=1.0,
        use_data_loss=True,
        use_pde_loss=True,
        use_constraint_loss=True,
        E_weight=1.0,
        B_weight=1.0,
        div_E_weight=1.0,
        div_B_weight=1.0,
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        use_weighted_mean=False
    ):
        super(MaxwellLoss, self).__init__()
        """初始化麦克斯韦方程损失函数
        
        参数:
            eps0: 真空介电常数
            mu0: 真空磁导率
            sigma: 电导率
            data_weight: 数据损失权重
            pde_weight: PDE损失权重
            constraint_weight: 约束损失权重
            use_data_loss: 是否使用数据损失
            use_pde_loss: 是否使用PDE损失
            use_constraint_loss: 是否使用约束损失
            E_weight: 电场权重
            B_weight: 磁场权重
            div_E_weight: 电场散度约束权重
            div_B_weight: 磁场散度约束权重
            Lx, Ly, Lz: 域尺寸
            use_weighted_mean: 是否使用加权平均
        """
        self.eps0 = eps0
        self.mu0 = mu0
        self.sigma = sigma
        self.data_weight = data_weight
        self.pde_weight = pde_weight
        self.constraint_weight = constraint_weight
        self.use_data_loss = use_data_loss
        self.use_pde_loss = use_pde_loss
        self.use_constraint_loss = use_constraint_loss
        self.E_weight = E_weight
        self.B_weight = B_weight
        self.div_E_weight = div_E_weight
        self.div_B_weight = div_B_weight
        self.Lx = Lx
        self.Ly = Ly
        self.Lz = Lz
        self.use_weighted_mean = use_weighted_mean
        
        # 初始化麦克斯韦方程
        self.maxwell_eq = MaxwellEquations(eps0, mu0, sigma)
        
        # 初始化损失函数
        self.lp_loss = LpLoss(p=2)
        
        if not self.use_data_loss:
            self.data_weight = 0
        if not self.use_pde_loss:
            self.pde_weight = 0
        if not self.use_constraint_loss:
            self.constraint_weight = 0
    
    def forward(self, pred, true, inputs=None, return_loss_dict=False):
        """计算麦克斯韦方程损失
        
        参数:
            pred: 预测值，形状为 [batch_size, x_dim, y_dim, z_dim, 6]
            true: 真实值，形状为 [batch_size, x_dim, y_dim, z_dim, 6]
            inputs: 输入值，可选
            return_loss_dict: 是否返回详细损失字典
            
        返回:
            总损失值，或(总损失值, 损失字典)
        """
        # 更新空间步长
        self.dx = self.Lx / (pred.shape[-4] - 1)
        self.dy = self.Ly / (pred.shape[-3] - 1)
        self.dz = self.Lz / (pred.shape[-2] - 1)
        
        if not return_loss_dict:
            loss = self.compute_loss(pred, true)
            return loss
        else:
            loss, loss_dict = self.compute_losses(pred, true)
            return loss, loss_dict
    
    def compute_loss(self, pred, true):
        """计算加权损失
        
        参数:
            pred: 预测值
            true: 真实值
            
        返回:
            总损失值
        """
        pred = pred.reshape(true.shape)
        
        # 提取电场和磁场
        E = pred[..., :3]
        B = pred[..., 3:]
        
        # 数据损失
        if self.use_data_loss:
            loss_data = self.data_loss(pred, true)
        else:
            loss_data = 0
        
        # PDE损失
        if self.use_pde_loss:
            DE, DB = self.maxwell_residual(E, B)
            loss_pde = self.pde_loss(DE, DB)
        else:
            loss_pde = 0
        
        # 约束损失
        if self.use_constraint_loss:
            div_E, div_B = self.maxwell_eq.compute_div_E(E, self.dx, self.dy, self.dz), \
                           self.maxwell_eq.compute_div_B(B, self.dx, self.dy, self.dz)
            loss_constraint = self.constraint_loss(div_E, div_B)
        else:
            loss_constraint = 0
        
        # 计算权重总和
        if self.use_weighted_mean:
            weight_sum = (
                self.data_weight
                + self.pde_weight
                + self.constraint_weight
            )
        else:
            weight_sum = 1.0
        
        # 总损失
        loss = (
            self.data_weight * loss_data
            + self.pde_weight * loss_pde
            + self.constraint_weight * loss_constraint
        ) / weight_sum
        
        return loss
    
    def compute_losses(self, pred, true):
        """计算加权损失和详细损失字典
        
        参数:
            pred: 预测值
            true: 真实值
            
        返回:
            总损失值和损失字典
        """
        pred = pred.reshape(true.shape)
        
        # 提取电场和磁场
        E = pred[..., :3]
        B = pred[..., 3:]
        
        loss_dict = {}
        
        # 数据损失
        if self.use_data_loss:
            loss_data, loss_E, loss_B = self.data_loss(
                pred, true, return_all_losses=True
            )
            loss_dict["loss_data"] = loss_data
            loss_dict["loss_E"] = loss_E
            loss_dict["loss_B"] = loss_B
        else:
            loss_data = 0
        
        # PDE损失
        if self.use_pde_loss:
            DE, DB = self.maxwell_residual(E, B)
            loss_pde, loss_DE, loss_DB = self.pde_loss(
                DE, DB, return_all_losses=True
            )
            loss_dict["loss_pde"] = loss_pde
            loss_dict["loss_DE"] = loss_DE
            loss_dict["loss_DB"] = loss_DB
        else:
            loss_pde = 0
        
        # 约束损失
        if self.use_constraint_loss:
            div_E = self.maxwell_eq.compute_div_E(E, self.dx, self.dy, self.dz)
            div_B = self.maxwell_eq.compute_div_B(B, self.dx, self.dy, self.dz)
            loss_constraint, loss_div_E, loss_div_B = self.constraint_loss(
                div_E, div_B, return_all_losses=True
            )
            loss_dict["loss_constraint"] = loss_constraint
            loss_dict["loss_div_E"] = loss_div_E
            loss_dict["loss_div_B"] = loss_div_B
        else:
            loss_constraint = 0
        
        # 计算权重总和
        if self.use_weighted_mean:
            weight_sum = (
                self.data_weight
                + self.pde_weight
                + self.constraint_weight
            )
        else:
            weight_sum = 1.0
        
        # 总损失
        loss = (
            self.data_weight * loss_data
            + self.pde_weight * loss_pde
            + self.constraint_weight * loss_constraint
        ) / weight_sum
        
        return loss, loss_dict
    
    def data_loss(self, pred, true, return_all_losses=False):
        """计算数据损失
        
        参数:
            pred: 预测值
            true: 真实值
            return_all_losses: 是否返回各分量损失
            
        返回:
            数据损失值，或(总损失, 各分量损失)
        """
        if return_all_losses:
            loss_E = self.lp_loss(pred[..., :3], true[..., :3])
            loss_B = self.lp_loss(pred[..., 3:], true[..., 3:])
            
            loss = self.E_weight * loss_E + self.B_weight * loss_B
            
            return loss, loss_E, loss_B
        else:
            loss = self.lp_loss(pred, true)
            return loss
    
    def maxwell_residual(self, E, B):
        """计算麦克斯韦方程残差
        
        参数:
            E: 电场
            B: 磁场
            
        返回:
            麦克斯韦方程残差
        """
        # 计算时间导数
        dE_dt = torch.gradient(E, dim=1)[0]  # 假设第一维是时间
        dB_dt = torch.gradient(B, dim=1)[0]
        
        # 计算法拉第定律残差
        dB_dt_pred = self.maxwell_eq.faraday_law(
            E, B, dx=self.dx, dy=self.dy, dz=self.dz
        )
        DB = dB_dt - dB_dt_pred
        
        # 计算安培定律残差
        dE_dt_pred = self.maxwell_eq.ampere_law(
            E, B, dx=self.dx, dy=self.dy, dz=self.dz
        )
        DE = dE_dt - dE_dt_pred
        
        return DE, DB
    
    def pde_loss(self, DE, DB, return_all_losses=False):
        """计算PDE损失
        
        参数:
            DE, DB: 麦克斯韦方程残差
            return_all_losses: 是否返回各分量损失
            
        返回:
            PDE损失值，或(总损失, 各分量损失)
        """
        if return_all_losses:
            loss_DE = self.lp_loss(DE, torch.zeros_like(DE))
            loss_DB = self.lp_loss(DB, torch.zeros_like(DB))
            
            loss = self.E_weight * loss_DE + self.B_weight * loss_DB
            
            return loss, loss_DE, loss_DB
        else:
            loss = (
                self.E_weight * self.lp_loss(DE, torch.zeros_like(DE))
                + self.B_weight * self.lp_loss(DB, torch.zeros_like(DB))
            )
            return loss
    
    def constraint_loss(self, div_E, div_B, return_all_losses=False):
        """计算约束损失
        
        参数:
            div_E: 电场散度
            div_B: 磁场散度
            return_all_losses: 是否返回各分量损失
            
        返回:
            约束损失值，或(总损失, 各分量损失)
        """
        if return_all_losses:
            loss_div_E = self.lp_loss(div_E, torch.zeros_like(div_E))
            loss_div_B = self.lp_loss(div_B, torch.zeros_like(div_B))
            
            loss = (
                self.div_E_weight * loss_div_E
                + self.div_B_weight * loss_div_B
            )
            
            return loss, loss_div_E, loss_div_B
        else:
            loss = (
                self.div_E_weight * self.lp_loss(div_E, torch.zeros_like(div_E))
                + self.div_B_weight * self.lp_loss(div_B, torch.zeros_like(div_B))
            )
            return loss